import numpy as np
import pandas as pd
import math
from typing import Optional, Tuple

# ========= 数据读取 =========
def load_series(df: pd.DataFrame, time_col: str, y_col: str) -> pd.DataFrame:
    """从 DataFrame 读取时间列与目标列，排序并清洗。"""
    s = df[[time_col, y_col]].copy()
    # 时间可为日期/数值；转失败则保留为原值
    try:
        s[time_col] = pd.to_datetime(s[time_col])
    except Exception:
        s[time_col] = pd.to_numeric(s[time_col], errors="coerce")
    s[y_col] = pd.to_numeric(s[y_col], errors="coerce")
    s = s.dropna().sort_values(time_col).reset_index(drop=True)
    s.columns = ["time", "y"]
    return s

# ========= GM(1,1)（等间隔）=========
class GM11:
    """
    经典灰色模型 GM(1,1)，适合等间隔序列。
    x0: 原序列；x1: 累加生成序列(AGO)
    模型：x0(k) + a*z1(k) = b，z1(k)=0.5*(x1(k)+x1(k-1))
    """
    def __init__(self):
        self.a_: Optional[float] = None
        self.b_: Optional[float] = None
        self.x0_: Optional[np.ndarray] = None

    @staticmethod
    def _ago(x: np.ndarray) -> np.ndarray:
        return np.cumsum(x)

    def fit(self, x0: np.ndarray):
        x0 = np.asarray(x0, dtype=float)
        if len(x0) < 4:
            raise ValueError("GM(1,1) 至少需要 4 个样本点")
        self.x0_ = x0
        x1 = self._ago(x0)
        z1 = -0.5 * (x1[1:] + x1[:-1])                  # 背景值取邻近均值
        B  = np.column_stack([z1, np.ones_like(z1)])    # [a, b]
        Y  = x0[1:]
        a, b = np.linalg.lstsq(B, Y, rcond=None)[0]
        self.a_, self.b_ = float(a), float(b)
        return self

    def fitted(self) -> np.ndarray:
        if self.a_ is None:
            raise RuntimeError("请先 fit")
        x0_1, a, b = self.x0_[0], self.a_, self.b_
        k = np.arange(1, len(self.x0_) + 1)
        x1_hat = (x0_1 - b/a) * np.exp(-a * (k - 1)) + b/a
        x0_hat = np.empty_like(self.x0_)
        x0_hat[0]  = self.x0_[0]             # 关键修正：首点等于原始首点
        x0_hat[1:] = np.diff(x1_hat)         # 其余用差分
        return x0_hat

    def predict(self, m: int) -> np.ndarray:
        if self.a_ is None:
            raise RuntimeError("请先 fit")
        x0_1, a, b = self.x0_[0], self.a_, self.b_
        k = np.arange(1, len(self.x0_) + m + 1)
        x1_hat = (x0_1 - b/a) * np.exp(-a * (k - 1)) + b/a
        x0_hat = np.diff(x1_hat, prepend=x1_hat[0])
        return x0_hat[len(self.x0_):]

# ========= 非等间隔 GM(1,1)（NIGM）=========
class NIGM11:
    """
    简化的非等间隔 GM(1,1)：
    连续型微分方程 dx/dt + a x = b
    区间解：x1(t_k) = (x1(t_{k-1}) - b/a) * exp(-a*Δt_k) + b/a
    用Δt加权最小二乘估计 a,b；背景值仍用邻近均值近似。
    """
    def __init__(self):
        self.a_: Optional[float] = None
        self.b_: Optional[float] = None
        self.t_: Optional[np.ndarray] = None
        self.x0_: Optional[np.ndarray] = None
        self.x1_: Optional[np.ndarray] = None

    @staticmethod
    def _ago(x: np.ndarray) -> np.ndarray:
        return np.cumsum(x)

    def fit(self, t: np.ndarray, x0: np.ndarray):
        t = np.asarray(t, dtype=float)
        x0 = np.asarray(x0, dtype=float)
        if len(x0) < 4:
            raise ValueError("NIGM(1,1) 至少需要 4 个样本点")

        # 时间轴平移到0，并做单位修正（若为秒级且跨度很大）
        t = t - t[0]
        if t.max() > 1e6:
            t = t / 86400.0  # 秒 -> 天
        self.t_, self.x0_ = t, x0
        x1 = self._ago(x0)
        self.x1_ = x1

        dt = np.diff(t)
        z  = -0.5 * (x1[1:] + x1[:-1])                # 背景值近似
        W  = np.diag(dt)                               # 用间隔加权
        B, Y = np.column_stack([z, np.ones_like(z)]), x0[1:]
        theta = np.linalg.solve(B.T @ W @ B, B.T @ W @ Y)
        self.a_, self.b_ = float(theta[0]), float(theta[1])
        return self

    def _x1_next(self, x1_prev: float, dt: float) -> float:
        a, b = self.a_, self.b_
        return (x1_prev - b/a) * math.exp(-a * dt) + b/a
    def fitted(self) -> np.ndarray:
        if self.a_ is None:
            raise RuntimeError("请先 fit")
        t = self.t_
        x1_hat = [self.x1_[0]]
        for k in range(1, len(t)):
            dt = t[k] - t[k-1]
            x1_hat.append(self._x1_next(x1_hat[-1], dt))
        x1_hat = np.array(x1_hat)
        x0_hat = np.empty_like(self.x0_)
        x0_hat[0]  = self.x0_[0]             # 关键修正
        x0_hat[1:] = np.diff(x1_hat)
        return x0_hat

    def predict(self, future_t: np.ndarray) -> np.ndarray:
        """future_t: 与历史同单位的未来时间戳数组（可以是 datetime 转秒/天后的数值）。"""
        if self.a_ is None:
            raise RuntimeError("请先 fit")
        t_hist = self.t_
        all_t  = np.concatenate([t_hist, np.asarray(future_t, dtype=float)])
        all_t  = all_t - all_t[0]
        if all_t.max() > 1e6:
            all_t = all_t / 86400.0

        # 先重建历史理论 x1，再外推
        x1_hat = [self.x1_[0]]
        for k in range(1, len(t_hist)):
            x1_hat.append(self._x1_next(x1_hat[-1], t_hist[k] - t_hist[k-1]))
        for k in range(len(t_hist), len(all_t)):
            x1_hat.append(self._x1_next(x1_hat[-1], all_t[k] - all_t[k-1]))
        x1_hat = np.array(x1_hat)
        x0_hat = np.diff(x1_hat, prepend=x1_hat[0])
        return x0_hat[len(t_hist):]
    
# ---- 基础函数：读并拼接所有 sheet ----
def load_all_sheets_as_series(xlsx_path: str, time_col: str, y_col: str) -> pd.DataFrame:
    """
    读取 Excel 的所有 sheet，提取 (time_col, y_col)，拼接成一个总体序列。
    返回 DataFrame: ['time','y','sheet']，按 time 升序且去重。
    """
    xls = pd.ExcelFile(xlsx_path)
    frames = []
    # print(xls.sheet_names)
    for sh in xls.sheet_names:
        df = pd.read_excel(xlsx_path, sheet_name=sh)
        # print(df.head())
        # print(time_col)
        # print(y_col)
        if time_col not in df.columns or y_col not in df.columns:
            # 某些 sheet 没这两列就跳过
            continue
        s = df[[time_col, y_col]].copy()
        # 统一类型
        try:
            s[time_col] = pd.to_datetime(s[time_col])
        except Exception:
            s[time_col] = pd.to_numeric(s[time_col], errors='coerce')
        s[y_col] = pd.to_numeric(s[y_col], errors='coerce')
        s = s.dropna().rename(columns={time_col:'time', y_col:'y'})
        s['sheet'] = sh
        frames.append(s)

    if not frames:
        raise ValueError("未在任何 sheet 中找到指定列。")

    all_s = pd.concat(frames, ignore_index=True)
    # 去除完全重复行（同一时间同一值）
    all_s = all_s.drop_duplicates(subset=['time','y']).sort_values('time').reset_index(drop=True)
    return all_s

# ========= 快速使用示例 =========
if __name__ == "__main__":
    # 1) 构造/读取数据
    df = load_all_sheets_as_series(
        "./P1data.xlsx",  # 请替换为实际的 Excel 文件路径
        time_col="时间",  # 时间列名
        y_col="整体就业率"     # 目标值列名
    )
    s = load_series(df, time_col="time", y_col="y")
    fut_2025 = pd.date_range("2025-01-01", "2025-12-31", periods=10)
    fut_2026 = pd.date_range("2026-01-01", "2026-12-31", periods=10)

    # 2) GM(1,1)：假定等间隔（直接按顺序取 y）
    gm = GM11().fit(s["y"].values)
    print("GM(1,1) in-sample MAPE=%.3f%%" % (
        100*np.mean(np.abs((s["y"].values - gm.fitted())/s["y"].values))
    ))
    gm_forecast = gm.predict(m=20)  # 预测后续20步（按等间隔步长）
    out = pd.DataFrame({"time": np.concatenate([fut_2025, fut_2026]), "y_pred": gm_forecast})
    out.to_excel("GM_outputs.xlsx", index=False)

    # 3) NIGM(1,1)：非等间隔（把 datetime 转为“秒”或“天”的数轴）
    t_num = s["time"].view("int64") / 1e9  # datetime64 -> 秒
    nigm  = NIGM11().fit(t_num, s["y"].values)
    # 未来时间点：以 2025/2026 各采 10 个点为例（日期 -> 数轴）
    
    future_t = np.concatenate([fut_2025.view("int64")/1e9, fut_2026.view("int64")/1e9])
    nigm_forecast = nigm.predict(future_t)

    # 4) 打印/保存
    print("NIGM(1,1) future points:", len(nigm_forecast))
    out = pd.DataFrame({"time": np.concatenate([fut_2025, fut_2026]), "y_pred": nigm_forecast})
    out.to_excel("NIGM_outputs.xlsx", index=False)
    
    
    
    import matplotlib.pyplot as plt

    # === 生成用于可视化的数据 ===
    # GM 轴：用整数索引
    idx_hist = np.arange(len(s))
    idx_fore = np.arange(len(s), len(s) + len(gm_forecast))

    gm_fit = gm.fitted()

    # NIGM 轴：用日期
    fut_times = np.concatenate([fut_2025, fut_2026])  # DatetimeIndex 拼接成一维数组

    # === 图1：GM(1,1) 拟合与预测（索引轴） ===
    plt.figure()
    plt.plot(idx_hist, s["y"].values, label="Observed (hist)")
    plt.plot(idx_hist, gm_fit, label="GM fitted")
    plt.plot(idx_fore, gm_forecast, label="GM forecast (20 steps)")
    plt.title("GM(1,1) — Fitted & Forecast (index-based)")
    plt.xlabel("Step index")
    plt.ylabel("Employment rate")
    plt.legend()
    plt.tight_layout()
    plt.savefig("GM_curve.png", dpi=150)

    # === 图2：NIGM(1,1) 未来预测（时间轴）===
    plt.figure()
    plt.plot(s["time"].values, s["y"].values, label="Observed (hist)")
    plt.plot(fut_times, nigm_forecast, label="NIGM forecast (2025–2026)")
    plt.title("NIGM(1,1) — Forecast on real time axis")
    plt.xlabel("Date")
    plt.ylabel("Employment rate")
    plt.legend()
    plt.tight_layout()
    plt.savefig("NIGM_curve.png", dpi=150)

    print("✅ 已保存图像：GM_curve.png, NIGM_curve.png 以及表格 GM_outputs.xlsx")
